from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Tuple

import httpx

from settings import RAGSettings


class OptimizationClientError(RuntimeError):
    """Raised when an optimization HTTP request fails."""


@dataclass
class DatasetItem:
    """Represent a single evaluation sample expected by diting."""

    id: str
    user_input: str
    expected_output: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None

    def as_dict(self) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "id": self.id,
            "user_input": self.user_input,
        }
        if self.expected_output is not None:
            payload["expected_output"] = self.expected_output
        if self.metadata is not None:
            payload["metadata"] = self.metadata
        return payload


@dataclass
class StageConfig:
    """Describe an optimization stage and its search space."""

    name: str
    parameters: Dict[str, Any] = field(default_factory=dict)
    similarity_threshold_range: Optional[Tuple[float, float]] = None
    similarity_threshold_step: Optional[float] = None
    context_recall_tokens_range: Optional[Tuple[int, int]] = None
    context_recall_tokens_step: Optional[int] = None
    extra_overrides: Dict[str, Any] = field(default_factory=dict)

    def as_dict(self) -> Dict[str, Any]:
        overrides: Dict[str, Any] = dict(self.extra_overrides)
        if self.similarity_threshold_range is not None:
            overrides["similarity_threshold_range"] = list(self.similarity_threshold_range)
        if self.similarity_threshold_step is not None:
            overrides["similarity_threshold_step"] = self.similarity_threshold_step
        if self.context_recall_tokens_range is not None:
            overrides["context_recall_tokens_range"] = list(self.context_recall_tokens_range)
        if self.context_recall_tokens_step is not None:
            overrides["context_recall_tokens_step"] = self.context_recall_tokens_step

        return {
            "name": self.name,
            "parameters": self.parameters,
            "parameters_overrides": overrides,
        }


def build_callback_config(
    *,
    stage_result_webhook: Optional[str] = None,
    result_webhook: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """Compose callback configuration payload accepted by diting."""
    callbacks: Dict[str, Any] = {}
    if stage_result_webhook:
        callbacks["stage_result_webhook"] = stage_result_webhook
    if result_webhook:
        callbacks["result_webhook"] = result_webhook
    return callbacks or None


def build_optimization_payload(
    *,
    schema_id: str,
    stage_sequence: Iterable[str],
    stages: Iterable[StageConfig],
    dataset_items: Iterable[DatasetItem],
    run_id: Optional[str] = None,
    metric: Optional[Dict[str, Any]] = None,
    global_context: Optional[Dict[str, Any]] = None,
    callback: Optional[Dict[str, Any]] = None,
    extra: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Utility to craft the request body diting expects."""
    payload: Dict[str, Any] = {
        "schema_id": schema_id,
        "stage_sequence": list(stage_sequence),
        "stages": [stage.as_dict() for stage in stages],
        "dataset_items": [item.as_dict() for item in dataset_items],
    }
    if run_id:
        payload["run_id"] = run_id
    if metric:
        payload["metric"] = metric
    if global_context:
        payload["global_context"] = global_context
    if callback:
        payload["callback"] = callback
    if extra:
        payload.update(extra)
    return payload


class DitingOptimizationClient:
    """HTTP client wrapper to interact with diting's optimization service."""

    def __init__(
        self,
        settings: RAGSettings,
        *,
        base_url: Optional[str] = None,
        token: Optional[str] = None,
        timeout: Optional[float] = None,
        client: Optional[httpx.Client] = None,
    ) -> None:
        self._settings = settings
        self._base_url = (base_url or settings.optimization_base_url or "").rstrip("/")
        self._token = token or settings.optimization_token
        self._timeout = timeout or settings.optimization_timeout
        self._client = client

    def _ensure_client(self) -> httpx.Client:
        if self._client is not None:
            return self._client
        return httpx.Client(timeout=self._timeout)

    def optimize_rag(self, payload: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """Execute optimization request and return best parameters with full response."""
        if not self._base_url:
            raise OptimizationClientError(
                "Optimization base URL not configured. Set RAG_OPTIMIZATION_BASE_URL."
            )

        url = f"{self._base_url}/api/v1/optimizations/rag"
        headers: Dict[str, str] = {}
        if self._token:
            headers["Authorization"] = f"Bearer {self._token}"

        client = self._ensure_client()
        owns_client = client is not self._client

        try:
            response = client.post(url, json=payload, headers=headers)
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            detail = exc.response.text
            raise OptimizationClientError(
                f"Optimization request failed with status {exc.response.status_code}: {detail}"
            ) from exc
        except httpx.HTTPError as exc:
            raise OptimizationClientError(f"Optimization request error: {exc}") from exc
        finally:
            if owns_client:
                client.close()

        body: Dict[str, Any] = response.json()
        best_parameters = body.get("best_parameters") or {}
        return best_parameters, body

